#!/usr/bin/env bash
# Created by weebney in 2023
# Licensed BSD-2-Clause

version="1.2.0"

camera=""
device_number=0
gphoto_args=(autofocusdrive=1)
ffmpeg_args=(-vcodec rawvideo -pix_fmt yuv420p -threads 0)
gphoto_errors=$(mktemp)
ffmpeg_errors=$(mktemp)
log_level="INFO"

# Trap
cleanup() {
	kill "$check_pid" &>/dev/null
	kill "$pipeline_pid" &>/dev/null
	rm -f "$gphoto_errors" "$ffmpeg_errors" &>/dev/null
}
trap cleanup EXIT

usage() {
	printf "Usage: webcamize [OPTIONS...]\n"
	printf "	-v, --version			Print version info and quit\n"
	printf "	-c, --camera NAME		Specify a gphoto2 camera to use; autodetects by default\n"
	printf "	-d, --device NUMBER		Specify the /dev/video device number to use (default: %s)\n" "${device_number}"
	printf "	-g, --gphoto-args ARGS		Pass arguments to gphoto2 (default: \"%s\")\n" "${gphoto_args[*]}"
	printf "	-f, --ffmpeg-args ARGS		Pass arguments to ffmpeg (default: \"%s\")\n" "${ffmpeg_args[*]}"
	printf "	-l, --log-level LEVEL		Set the log level (INFO, WARN, FATAL; default: INFO)\n"
	printf "	-h, --help			Show this help message\n"
}
warn() {
	[ "$log_level" != "FATAL" ] || return
	# shellcheck disable=SC2059
	printf "\e[33m[WARN]\e[0m $1\n" "${@:2}" >&2
}
info() {
	[ "$log_level" = "INFO" ] || return
	# shellcheck disable=SC2059
	printf "\e[32m[INFO]\e[0m $1\n" "${@:2}" >&2
}
fatal() {
	# shellcheck disable=SC2059
	printf "\e[31m[FATAL]\e[0m $1\n" "${@:2}"
	exit 1
}

# Command line options
missing_arg() {
	fatal "Missing argument for %s" "$1"
}
while (("$#")); do
	case "$1" in
	-v | --version)
		info "Using webcamize %s" "${version}"
		exit
		;;
	-d | --device)
		if [ -n "$2" ]; then
			# Validate input
			if ! [[ "$2" =~ ^[0-9]+$ ]] || (("$2" < 0)); then
				fatal "Argument for %s must be a non-negative integer" "$1"
			fi
			device_number=$2
			shift 2
		else
			missing_arg "$1"
		fi
		;;
	-c | --camera)
		if [ -n "$2" ]; then
			camera="--camera \"$2\""
			shift 2
		else
			missing_arg "$1"
		fi
		;;
	-g | --gphoto-args)
		if [ -n "$2" ]; then
			IFS=' ' read -ra gphoto_args <<<"$2"
			shift 2
		else
			missing_arg "$1"
		fi
		;;
	-f | --ffmpeg-args)
		if [ -n "$2" ]; then
			IFS=' ' read -ra ffmpeg_args <<<"$2"
			shift 2
		else
			missing_arg "$1"
		fi
		;;
	-l | --log-level)
		if [ -n "$2" ]; then
			case "$2" in
			INFO | WARN | FATAL)
				log_level="$2"
				;;
			*)
				fatal "Invalid log level. Valid options are: INFO, WARN, FATAL"
				;;
			esac
			shift 2
		else
			missing_arg "$1"
		fi
		;;
	-h | --help)
		usage
		exit
		;;
	--)
		shift
		break
		;;
	-*)
		fatal "Unsupported flag %s" "$1"
		;;
	*)
		shift
		;;
	esac
done

# Verify dependencies
for cmd in gphoto2 ffmpeg sudo pgrep modinfo modprobe lsmod; do
	if ! command -v $cmd &>/dev/null; then
		fatal "Dependency %s could not be found" "$cmd"
	fi
done

# Verify that gphoto2 is working
autodetect=$(gphoto2 --auto-detect | awk 'NR==3' | sed 's/  .*//')
if [ -z "${autodetect}" ]; then
	fatal "Couldn't detect any cameras with gphoto2"
fi

# Set camera name
if [ -z "${camera}" ]; then
	camera_name="${autodetect}"
else
	camera_name=${camera#* }
fi

# Verify that the v4l2loopback module is installed
if ! modinfo v4l2loopback &>/dev/null; then
	warn "Possibly missing v4l2loopback module"
fi

# Verify that the loopback device exists
if [ ! -e /dev/video"${device_number}" ]; then
	warn "Could not find /dev/video%s" "${device_number}"
	
	# Unload v4l2loopback so we can add new devices
	if lsmod | grep -q "v4l2loopback"; then
		warn "Reloading v4l2loopback; this will disable in-use camera feeds"
		if ! sudo modprobe -r v4l2loopback; then
			fatal "Error unloading v4l2loopback module"
		fi
	fi

	info "Attempting to add /dev/video%s via modprobe" "${device_number}"
	
	# Add the required device
	if ! sudo modprobe v4l2loopback video_nr="${device_number}" card_label="${camera_name} (webcamize)" exclusive_caps=1; then
		fatal "Error loading v4l2loopback module"
	fi

	# Verify that the device exists
	if [ ! -e /dev/video"${device_number}" ]; then
		fatal "Failed to added /dev/video%s" "${device_number}"
	else
		info "Successfully added /dev/video%s!" "${device_number}"
	fi
fi

info "Starting %s on /dev/video%s" "${camera_name}" "${device_number}"

# Main execution below
(gphoto2 "${gphoto_args[@]}" --capture-movie --stdout 2>"$gphoto_errors" |
	ffmpeg -i - "${ffmpeg_args[@]}" -f v4l2 /dev/video"${device_number}" -hide_banner -loglevel warning -nostats >/dev/null 2>"$ffmpeg_errors") &
pipeline_pid=$!

check_video_output() {
	prev_output=$(head -c 64 /dev/video"${device_number}" 2>/dev/null)
	while true; do
		if ! kill -0 "$pipeline_pid" &>/dev/null; then
			return
		fi
		current_output=$(head -c 64 /dev/video"${device_number}" 2>/dev/null)
		if [ "$current_output" != "$prev_output" ]; then
			info "Successfully started!"
			break
		fi
		prev_output=$current_output
	done
}
check_video_output &
check_pid=$!

# Manage pipeline errors
wait $pipeline_pid
pipeline_status=$?
if [ $pipeline_status -ne 0 ]; then
	warn "Pipeline exited with non-0 status (%s)" "${pipeline_status}"
	sed -i '0,/abort./d' "${gphoto_errors}"
	if [ -s "$gphoto_errors" ]; then
		fatal "gphoto2: %s" "$(cat "$gphoto_errors")" &
	fi
	if [ -s "$ffmpeg_errors" ]; then
		fatal "ffmpeg: %s" "$(cat "$ffmpeg_errors")"
	fi
fi
